from torch.autograd import Variable
import torch.nn as nn


class Discriminator(nn.Module):

    def __init__(self, num_classes, ndf=64):
        super(Discriminator, self).__init__()

        self.conv1 = nn.Conv2d(num_classes + 3, ndf, 4, 2, 1)
        self.conv2 = nn.Conv2d(ndf, ndf * 2, 4, 2, 1)
        self.conv3 = nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1)
        self.conv4 = nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1)
        self.pool = nn.AvgPool2d((32, 32))
        self.fc = nn.Linear(ndf * 8, 1)
        self.leaky_relu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        self.drop = nn.Dropout2d(0.5)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.conv1(x)
        x = self.leaky_relu(x)
        x = self.drop(x)

        x = self.conv2(x)
        x = self.leaky_relu(x)
        x = self.drop(x)

        x = self.conv3(x)
        x = self.leaky_relu(x)
        x = self.drop(x)

        x = self.conv4(x)
        x = self.leaky_relu(x)

        maps = self.pool(x)
        out = maps.view(maps.size(0), -1)
        out = self.fc(out)

        return out, maps
